"""
augmentation 数据增强操作工具包
"""
import numpy as np
import cv2

def t_random(min=0, max=1):
    """生成指定范围内的均匀分布随机浮点数
    Args:
        min (float): 最小值，默认为0
        max (float): 最大值，默认为1
    Returns:
        float: [min, max)区间内的随机数
    """
    return min + (max - min) * np.random.rand()


def t_randint(min, max):
    """生成指定范围内的随机整数
    Args:
        min (int): 最小值（包含）
        max (int): 最大值（不包含）
    Returns:
        int: [min, max)区间内的随机整数
    """
    return np.random.randint(low=min, high=max)

class augCompose(object):
    """数据增强组合器，按顺序执行多个增强操作
    Attributes:
        transforms (list): 包含(操作函数, 执行概率)元组的列表
    """
    def __init__(self, transforms=None):
        """初始化增强器
        Args:
            transforms (list, optional): 增强操作列表，格式为[(op, prob),...]
        """
        self.transforms = transforms

    def __call__(self, img, mask):
        """执行增强操作链
        Args:
            img (numpy.ndarray): 原始图像
            mask (numpy.ndarray): 对应的掩膜图像
        Returns:
            tuple: 增强后的图像和掩膜
        """
        if self.transforms is not None:
            for op, prob in self.transforms:
                if t_random() <= prob:  # 按概率执行操作
                    img, mask = op(img, mask)
        return img, mask

def RandomFlip(img, mask, FLIP_LEFT_RIGHT=True, FLIP_TOP_BOTTOM=True):
    """随机执行水平和/或垂直翻转
    Args:
        img (numpy.ndarray): 输入图像
        mask (numpy.ndarray): 输入掩膜
        FLIP_LEFT_RIGHT (bool): 是否启用水平翻转
        FLIP_TOP_BOTTOM (bool): 是否启用垂直翻转
    Returns:
        tuple: 翻转后的图像和掩膜
    """
    if FLIP_LEFT_RIGHT and t_random() < 0.5:  # 50%概率水平翻转
        img = cv2.flip(img, 1)  # 1表示水平翻转
        mask = cv2.flip(mask, 1)

    if FLIP_TOP_BOTTOM and t_random() < 0.5:  # 50%概率垂直翻转
        img = cv2.flip(img, 0)  # 0表示垂直翻转
        mask = cv2.flip(mask, 0)
    return img, mask


def RandomBlur(img, mask):
    """随机应用不同类型模糊
    执行概率：高斯模糊(20%) > 均值模糊(15%) > 中值模糊(10%)
    Args:
        img (numpy.ndarray): 输入图像
        mask (numpy.ndarray): 输入掩膜
    Returns:
        tuple: 模糊处理后的图像和掩膜
    """
    r = 5  # 卷积核半径

    if t_random() < 0.2:
        return cv2.GaussianBlur(img,(r,r),0), mask  # 高斯模糊

    if t_random() < 0.15:
        return cv2.blur(img,(r,r)), mask  # 均值模糊

    if t_random() < 0.1:
        return cv2.medianBlur(img,r), mask  # 中值模糊

    return img, mask

def RandomColorJitter(img, mask, brightness=32, contrast=0.5, saturation=0.5, hue=0.1,
                        prob=0.5):
    """随机颜色扰动（亮度/对比度/饱和度/色调）
    Args:
        img (numpy.ndarray): 输入图像（BGR格式）
        mask (numpy.ndarray): 输入掩膜
        brightness (int): 亮度变化幅度
        contrast (float): 对比度变化范围（0.5表示±50%）
        saturation (float): 饱和度变化范围
        hue (float): 色调变化范围（0-0.5）
        prob (float): 各操作独立执行概率
    Returns:
        tuple: 调整后的图像和掩膜
    """
    if brightness != 0 and t_random() > prob:
        img = _Brightness(img, delta=brightness)
    if contrast != 0 and t_random() > prob:
        img = _Contrast(img, var=contrast)
    if saturation != 0 and t_random() > prob:
        img = _Saturation(img, var=saturation)
    if hue != 0 and t_random() > prob:
        img = _Hue(img, var=hue)

    return img, mask


def _Brightness(img, delta=32):
    """调整亮度（内部函数）
    Args:
        img (numpy.ndarray): 输入图像
        delta (int): 亮度变化幅度（±delta）
    Returns:
        numpy.ndarray: 调整后的图像
    """
    img = img.astype(np.float32) + t_random(-delta, delta)  # 随机亮度偏移
    img = np.clip(img, 0, 255)  # 数值截断
    return img.astype(np.uint8)


def _Contrast(img, var=0.3):
    """调整对比度（内部函数）
    Args:
        img (numpy.ndarray): 输入图像
        var (float): 对比度变化范围（0.7-1.3）
    Returns:
        numpy.ndarray: 调整后的图像
    """
    gs = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY).mean()  # 计算灰度均值
    alpha = 1.0 + t_random(-var, var)  # 对比度系数
    img = alpha * img.astype(np.float32) + (1 - alpha) * gs  # 线性插值
    img = np.clip(img, 0, 255)
    return img.astype(np.uint8)


def _Hue(img, var=0.05):
    """调整色调（内部函数）
    Args:
        img (numpy.ndarray): 输入图像
        var (float): 色调变化范围（0-0.5）
    Returns:
        numpy.ndarray: 调整后的图像
    """
    var = t_random(-var, var)
    # 随机选择RGB或BGR转换方式
    to_HSV, from_HSV = [
        (cv2.COLOR_RGB2HSV, cv2.COLOR_HSV2RGB),
        (cv2.COLOR_BGR2HSV, cv2.COLOR_HSV2BGR)][t_randint(0, 2)]
    hsv = cv2.cvtColor(img, to_HSV).astype(np.float32)

    hue = hsv[:, :, 0] / 179. + var  # OpenCV的H通道范围是0-179
    hue = hue - np.floor(hue)  # 归一化处理
    hsv[:, :, 0] = hue * 179.  # 恢复H通道范围

    img = cv2.cvtColor(hsv.astype('uint8'), from_HSV)
    return img


def _Saturation(img, var=0.3):
    """调整饱和度（内部函数）
    Args:
        img (numpy.ndarray): 输入图像
        var (float): 饱和度变化范围（0.7-1.3）
    Returns:
        numpy.ndarray: 调整后的图像
    """
    gs = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    gs = np.expand_dims(gs, axis=2)  # 增加通道维度
    alpha = 1.0 + t_random(-var, var)
    img = alpha * img.astype(np.float32) + (1 - alpha) * gs.astype(np.float32)
    img = np.clip(img, 0, 255)
    return img.astype(np.uint8)